import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

# 书上题目5.1 数据
def create_data():
    datasets = [['青年', '否', '否', '一般', '否'],
               ['青年', '否', '否', '好', '否'],
               ['青年', '是', '否', '好', '是'],
               ['青年', '是', '是', '一般', '是'],
               ['青年', '否', '否', '一般', '否'],
               ['中年', '否', '否', '一般', '否'],
               ['中年', '否', '否', '好', '否'],
               ['中年', '是', '是', '好', '是'],
               ['中年', '否', '是', '非常好', '是'],
               ['中年', '否', '是', '非常好', '是'],
               ['老年', '否', '是', '非常好', '是'],
               ['老年', '否', '是', '好', '是'],
               ['老年', '是', '否', '好', '是'],
               ['老年', '是', '否', '非常好', '是'],
               ['老年', '否', '否', '一般', '否'],
               ]
    labels = [u'年龄', u'有工作', u'有自己的房子', u'信贷情况', u'类别']

    x = np.array(datasets)
    y = x[:,-1]
    x = x[:,:-1]
    # 返回数据集和每个维度的名称
    return x, y, labels

# 节点定义
class Node(object):

    def __init__(self,dim,feature,label,child=None):
        self.dim = dim # 切分此节点的特征的索引,叶节点为-1
        self.feature = feature # 如果节点为叶节点则为-1.否则表示对应样本的特征,根节点为0
        self.label   = label   # 如果节点为叶节点则表示对应样本的分类，否则为-1
        self.child   = child

    def __str__(self):
        return str(self.label)  #print 一个 Node 类时会打印 __str__ 的返回值

class DTree(): # 决策树 ID3 算法

    def __init__(self,x,y,e=0.2):
        self.e = e
        self.root = self.fit(x,y)

    def fit(self,x,y,feature=0):

        m,_ = x.shape

        labels_num = len(np.unique(y))
        entrpy = self.calc_ent(y)

        if labels_num == 0:
            return None

        elif labels_num == 1:
            return Node(-1,feature,y[0])

        elif entrpy < self.e:
            label_counts = {}
            for i in y:
                if i not in label_counts.keys():
                    label_counts[i] = 0
                label_counts[i] += 1
            index = np.argmax(label_counts.values())
            return Node(-1,feature,label_counts.keys()[index])

        else:

            dim = self.calc_gain(x,y)

            x_cut = {}
            for i in range(m):
                if x[i,dim] not in x_cut.keys():
                    x_cut[x[i,dim]] = []
                    x_cut[x[i,dim]].append([])
                    x_cut[x[i,dim]].append([])
                x_cut[x[i,dim]][0].append(x[i,:])
                x_cut[x[i,dim]][1].append(y[i])
            child = []
            for j,i in x_cut.items():
                child.append( self.fit(np.array(i[0]),np.array(i[1]),j) )

            return Node(dim,feature,-1,child)

    def score(self,x_test,y_test):

        y_pred = []
        for i in range(len(y_test)):
            y_pred.append(self.predict( x_test[i,:] ))

        total_num = len(y_test)
        corre_num = np.sum( (y_pred == y_test).astype(int) )
        acc = corre_num / total_num

        return acc

    def predict(self,x):
        label = self.serach(x)
        return label

    def calc_ent(self,y): # 计算熵

        n = len(y)
        label_counts = {}
        for i in y:
            if i not in label_counts.keys():
                label_counts[i] = 0
            label_counts[i] += 1
        p = np.array(list(label_counts.values())) / n
        ent = - np.sum( p * np.log2(p) )

        return ent

    def cond_ent(self,x,y,axis=0): # 计算条件熵

        n = len(y)
        x_counts = {}
        for i in range(n):
            if x[i,axis] not in x_counts.keys():
                x_counts[x[i,axis]] = []
                x_counts[x[i,axis]].append(0)
                x_counts[x[i,axis]].append([])
            x_counts[x[i,axis]][0] += 1
            x_counts[x[i,axis]][1].append(y[i])
        ent = np.sum([ p[0]/n * self.calc_ent(p[1]) for p in x_counts.values() ])

        return ent

    def calc_gain(self,x,y): # 计算互信息
        
        ent = self.calc_ent(y)
        cond_e = []
        _,n = x.shape
        for i in range(n):
            cond_e.append( self.cond_ent(x,y,i) )
        gain = ent - np.array(cond_e) 
        dim = np.argmax(gain)

        return dim

    def serach(self,x): # 搜索
        Node = self.root
        while Node.child != None:
            for i in Node.child:
                if x[Node.dim] == i.feature:
                    Node = i
                    break
        return Node.label

# 主函数
def main():

    # 加载数据，打印基本信息
    x, y, _ = create_data()

    # 创建与训练
    DecisionTree = DTree(x,y)

    # 测试
    score = DecisionTree.score(x,y)
    print(score)

if __name__ == '__main__':
    main()